# -*- coding: UTF-8 -*-
"""
@项目名称：control_gui.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-12-14:04
"""
from rabbit_mq.rmq_config import logger, cfg
import gradio as gr
import argparse
from util_tools.run_main import run_main, get_task_progress

parse = argparse.ArgumentParser()
parse.add_argument("--ip", type=str, default="0.0.0.0")
parse.add_argument("--port", type=int, default=50000)
args = parse.parse_args()
logger.info(f'服务地址：{args.ip}:{args.port}')
model_list = cfg['model'].keys()

with gr.Blocks(title='AIDD_Control_Generate', theme=None) as demo:
    with gr.Tab("分子生成"):
        gr.Markdown(
            """
            edu_cqu_ADO（3个关键文件）： </br>
            10A口袋、该蛋白质具备生长点的配体sdf，该蛋白质具备生长点配体sdf的smiles </br>
            model_TG： </br>
            经验证：仅能使用https://www.rcsb.org下载cif格式蛋白质，PyMOL、薛定谔导处cif无法使用 </br>
            """)
        with gr.Row():
            model_name_list = gr.CheckboxGroup(model_list, label="模型列表", info="选择你的对接模型")
        with gr.Row():
            with gr.Column(scale=1):
                protein_file = gr.File(file_count="single", label="包含配体的蛋白质文件", file_types=[".pdb", ".cif"])
            with gr.Column(scale=1):
                smiles_file = gr.File(file_count="single", label="阳性分子的SMILES文件csv,必须包含【SMILES】字段", file_types=[".csv"])
        with gr.Row():
            with gr.Column(scale=1):
                cqu_pocket_10A_file = gr.File(file_count="single", label="重大：10A口袋.pdb，", file_types=[".pdb"])
            with gr.Column(scale=1):
                cqu_ado_sdf_file = gr.File(file_count="single", label="重大：具备生长点的sdf构象，与smiles所属同一个化合物", file_types=[".sdf"])
            with gr.Column(scale=1):
                cqu_ado_smi_file = gr.File(file_count="single", label="重大：具备生长点的smiles，与sdf所属同一个化合物", file_types=[".smi"])

        with gr.Row():
            with gr.Column(scale=1):
                run_button = gr.Button("开始生成")
        with gr.Row():
            result_path = gr.Text(label='结果文件路径', info='')
            # 添加示例到界面中
        inputs = [model_name_list, protein_file, smiles_file, cqu_pocket_10A_file, cqu_ado_sdf_file, cqu_ado_smi_file]
        outputs = [result_path]
        run_button.click(run_main, inputs=inputs, outputs=outputs)
        examples = gr.Examples(
            examples=[
                [['edu_cqu_ADO', 'model_TG'],  # 'model_PG', 'model_SG', 'model_MC',
                 f"{cfg['generate_input_path']}/4bvw/4bvw.cif",
                 f"{cfg['generate_input_path']}/4bvw/ligand_positive.csv",
                 f"{cfg['generate_input_path']}/eud_cqu_ado/pocket_10A.pdb",
                 f"{cfg['generate_input_path']}/eud_cqu_ado/scaffold_good_c.sdf",
                 f"{cfg['generate_input_path']}/eud_cqu_ado/scaffold_c.smi"
                 ]
            ],
            inputs=[model_name_list, protein_file, smiles_file, cqu_pocket_10A_file, cqu_ado_sdf_file, cqu_ado_smi_file],
        )
    with gr.Tab("任务进度"):
        with gr.Row():
            with gr.Column(scale=1):
                task_id = gr.Text(label='任务ID', info='')
                task_progress = gr.Text(label='任务进度', info='')
            with gr.Column(scale=1):
                result_file_list = gr.File(file_count="multiple", label="")

        task_button = gr.Button("获取任务进度")
        task_button.click(get_task_progress, inputs=[task_id], outputs=[task_progress, result_file_list])

    demo.launch(server_name=args.ip, server_port=args.port, root_path=f"/AIDD_Control_Generate", favicon_path='images/基因统计.svg', share=False)  #
# conda run -n AIDD_Control python control_gui.py
